// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/metadata_filter.h>

#include <melon/container/f14_map.h>
#include <pollux/dwio/common/scan_spec.h>
#include <pollux/expression/expr_to_subfield_filter.h>

namespace kumo::pollux::common {
    namespace {
        using LeafResults =
        melon::F14FastMap<const MetadataFilter::LeafNode *, std::vector<uint64_t> *>;
    }

    struct MetadataFilter::Node {
        static std::unique_ptr<Node> fromExpression(
            const core::ITypedExpr &,
            core::ExpressionEvaluator *,
            bool negated);

        virtual ~Node() = default;

        virtual void addToScanSpec(ScanSpec &) const = 0;

        virtual uint64_t *eval(LeafResults &, int size) const = 0;

        virtual std::string toString() const = 0;
    };

    class MetadataFilter::LeafNode : public Node {
    public:
        LeafNode(Subfield &&field, std::unique_ptr<Filter> filter)
            : field_(std::move(field)), filter_(std::move(filter)) {
        }

        void addToScanSpec(ScanSpec &scanSpec) const override {
            scanSpec.getOrCreateChild(field_)->addMetadataFilter(this, filter_.get());
        }

        uint64_t *eval(LeafResults &leafResults, int) const override {
            if (auto it = leafResults.find(this); it != leafResults.end()) {
                return it->second->data();
            }
            return nullptr;
        }

        const Subfield &field() const {
            return field_;
        }

        std::string toString() const override {
            return field_.toString() + ":" + filter_->toString();
        }

    private:
        Subfield field_;
        std::unique_ptr<Filter> filter_;
    };

    struct MetadataFilter::AndNode : Node {
        static std::unique_ptr<Node> create(
            std::unique_ptr<Node> lhs,
            std::unique_ptr<Node> rhs) {
            if (!lhs) {
                return rhs;
            }
            if (!rhs) {
                return lhs;
            }
            return std::make_unique<AndNode>(std::move(lhs), std::move(rhs));
        }

        AndNode(std::unique_ptr<Node> lhs, std::unique_ptr<Node> rhs)
            : lhs_(std::move(lhs)), rhs_(std::move(rhs)) {
        }

        void addToScanSpec(ScanSpec &scanSpec) const override {
            lhs_->addToScanSpec(scanSpec);
            rhs_->addToScanSpec(scanSpec);
        }

        uint64_t *eval(LeafResults &leafResults, int size) const override {
            auto *l = lhs_->eval(leafResults, size);
            auto *r = rhs_->eval(leafResults, size);
            if (!l) {
                return r;
            }
            if (!r) {
                return l;
            }
            bits::orBits(l, r, 0, size);
            return l;
        }

        std::string toString() const override {
            return "and(" + lhs_->toString() + "," + rhs_->toString() + ")";
        }

    private:
        std::unique_ptr<Node> lhs_;
        std::unique_ptr<Node> rhs_;
    };

    struct MetadataFilter::OrNode : Node {
        static std::unique_ptr<Node> create(
            std::unique_ptr<Node> lhs,
            std::unique_ptr<Node> rhs) {
            if (!lhs || !rhs) {
                return nullptr;
            }
            return std::make_unique<OrNode>(std::move(lhs), std::move(rhs));
        }

        OrNode(std::unique_ptr<Node> lhs, std::unique_ptr<Node> rhs)
            : lhs_(std::move(lhs)), rhs_(std::move(rhs)) {
        }

        void addToScanSpec(ScanSpec &scanSpec) const override {
            lhs_->addToScanSpec(scanSpec);
            rhs_->addToScanSpec(scanSpec);
        }

        uint64_t *eval(LeafResults &leafResults, int size) const override {
            auto *l = lhs_->eval(leafResults, size);
            auto *r = rhs_->eval(leafResults, size);
            if (!l || !r) {
                return nullptr;
            }
            bits::andBits(l, r, 0, size);
            return l;
        }

        std::string toString() const override {
            return "or(" + lhs_->toString() + "," + rhs_->toString() + ")";
        }

    private:
        std::unique_ptr<Node> lhs_;
        std::unique_ptr<Node> rhs_;
    };

    namespace {
        const core::CallTypedExpr *asCall(const core::ITypedExpr *expr) {
            return dynamic_cast<const core::CallTypedExpr *>(expr);
        }
    } // namespace

    std::unique_ptr<MetadataFilter::Node> MetadataFilter::Node::fromExpression(
        const core::ITypedExpr &expr,
        core::ExpressionEvaluator *evaluator,
        bool negated) {
        auto *call = asCall(&expr);
        if (!call) {
            return nullptr;
        }
        if (call->name() == "and") {
            auto lhs = fromExpression(*call->inputs()[0], evaluator, negated);
            auto rhs = fromExpression(*call->inputs()[1], evaluator, negated);
            return negated
                       ? OrNode::create(std::move(lhs), std::move(rhs))
                       : AndNode::create(std::move(lhs), std::move(rhs));
        }
        if (call->name() == "or") {
            auto lhs = fromExpression(*call->inputs()[0], evaluator, negated);
            auto rhs = fromExpression(*call->inputs()[1], evaluator, negated);
            return negated
                       ? AndNode::create(std::move(lhs), std::move(rhs))
                       : OrNode::create(std::move(lhs), std::move(rhs));
        }
        if (call->name() == "not") {
            return fromExpression(*call->inputs()[0], evaluator, !negated);
        }
        try {
            Subfield subfield;
            auto filter =
                    exec::ExprToSubfieldFilterParser::getInstance()
                    ->leafCallToSubfieldFilter(*call, subfield, evaluator, negated);
            if (!filter) {
                return nullptr;
            }
            POLLUX_CHECK(
                subfield.valid(),
                "Invalid subfield from expression: {}",
                expr.toString());
            return std::make_unique<LeafNode>(std::move(subfield), std::move(filter));
        } catch (const PolluxException &) {
            KLOG(WARNING) << "Fail to convert expression to metadata filter: "
                 << expr.toString();
            return nullptr;
        }
    }

    MetadataFilter::MetadataFilter(
        ScanSpec &scanSpec,
        const core::ITypedExpr &expr,
        core::ExpressionEvaluator *evaluator)
        : root_(Node::fromExpression(expr, evaluator, false)) {
        if (root_) {
            root_->addToScanSpec(scanSpec);
        }
    }

    void MetadataFilter::eval(
        std::vector<std::pair<const LeafNode *, std::vector<uint64_t> > > &
        leafNodeResults,
        std::vector<uint64_t> &finalResult) {
        if (!root_) {
            return;
        }

        LeafResults leafResults;
        for (auto &[leaf, result]: leafNodeResults) {
            POLLUX_CHECK_EQ(
                result.size(),
                finalResult.size(),
                "Result size mismatch: {}",
                leaf->field().toString());
            POLLUX_CHECK(
                leafResults.emplace(leaf, &result).second,
                "Duplicate results: {}",
                leaf->field().toString());
        }
        const auto bitCount = finalResult.size() * 64;
        if (auto *combined = root_->eval(leafResults, bitCount)) {
            bits::orBits(finalResult.data(), combined, 0, bitCount);
        }
    }

    std::string MetadataFilter::toString() const {
        return !root_ ? "" : root_->toString();
    }
} // namespace kumo::pollux::common
